from typing import List

from pymilvus import Collection, connections, db, FieldSchema, DataType, CollectionSchema, utility
from tabulate import tabulate

# 加载 Embedding 模型
from models.embedding.load_embedding_model import LoadEmbeddingModel

model = LoadEmbeddingModel()
model.load()


# 创建向量数据库Collection
# text_data 业务数据
def create_milvus(text_data: str):
    # 连接到 Milvus 服务器
    connections.connect(host="172.16.176.66", port=19530)
    # 使用数据库
    db.using_database("default")

    # ============================================================

    # 创建主键
    pk_id = FieldSchema(
        name="pk_id",
        dtype=DataType.INT64,
        is_primary=True,
    )
    # 创建标量，用来记录视频时间
    video_time = FieldSchema(
        name="video_time",
        dtype=DataType.VARCHAR,
        max_length=200,
        default_value="Unknown"
    )
    # 创建向量，用来记录视频中的旁白文字
    video_context = FieldSchema(
        name="video_context",
        dtype=DataType.VARCHAR,
        max_length=500,
        default_value="Unknown"
    )
    # 创建向量，用来记录视频中的旁白文字向量化
    video_context_vector = FieldSchema(
        name="video_context_vector",
        dtype=DataType.FLOAT_VECTOR,
        dim=1024
    )
    # 将列名放入CollectionSchema对象中，为创建Collection做准备
    schema = CollectionSchema(
        fields=[pk_id, video_time, video_context, video_context_vector],
        description="根据旁白文字查询视频时间",
        enable_dynamic_field=True
    )
    # Collection名
    collection_name = "video"

    # 在default数据库中创建Collection
    collection = Collection(
        name=collection_name,
        schema=schema,
        using='default'
    )

    # ============================================================

    # 创建向量索引
    index_params = {
        "metric_type": "L2",
        "index_type": "HNSW",
        "params": {
            "M": 8,
            "efConstruction": 64
        }
    }

    # 为Collection中的向量列名video_context_vector添加向量索引
    collection.create_index(
        field_name="video_context_vector",
        index_params=index_params
    )
    utility.index_building_progress("video")

    # ============================================================

    # 插入业务数据
    data = []

    # 将文本数据分割成每个时间段的旁白文字
    segments = [segment.strip() for segment in text_data.split("\n\n")]
    # 遍历每个时间段的旁白文字
    for idx, segment in enumerate(segments):
        # ['时间范围', '旁白文字']
        temp = segment.split("\n")
        print(temp)
        # 获取时间范围
        time_range = temp[0]
        # 获取旁白文字
        text = temp[1]
        # 使用嵌入模型获取旁白文字的向量
        embedding_vector = model.encode(text)
        # 创建数据点字典
        data_point = {
            "pk_id": idx + 1,  # 主键
            "video_time": time_range,  # 视频时间
            "video_context": text,  # 旁白文字
            "video_context_vector": embedding_vector.tolist()  # 旁白文字向量
        }
        # 添加到数据列表
        data.append(data_point)

    # 插入数据到Milvus
    collection.insert(data)
    # 加载Collection
    collection.load()


# 查询向量数据库
# query_vector 查询条件向量
def search_milvus(query_vector: List):
    # 连接到 Milvus 服务器
    conn = connections.connect(host="172.16.176.66", port=19530)
    # 使用数据库
    db.using_database("default")

    # 获取现有的集合
    collection = Collection("video")
    # 加载集合到内存
    collection.load()

    # 准备搜索参数
    search_params = {
        "metric_type": "L2",
        "offset": 0,
        "params": {"ef": 250}
    }

    # 进行搜索操作
    results = collection.search(
        data=[query_vector],  # 待搜索的数据，这里使用二维数组表示
        anns_field="video_context_vector",  # 指定进行向量搜索的列名
        # `param` 中 `offset` 和 `limit` 的总和应小于 16384
        param=search_params,  # 搜索参数，包括偏移和限制等
        limit=1,  # 搜索结果返回的数量限制
        expr=None,  # 表达式，用于进一步筛选结果，这里设置为 None
        # 指定要从搜索结果中检索的列名(非内置列名都要写)
        output_fields=['pk_id', 'video_time', 'video_context'],
        consistency_level="Strong"  # 一致性级别，这里设置为“Strong”
    )

    # 输出结果格式化成横向表格
    # 定义列头
    headers = ["pk_id", "score", "video_time", "video_context"]
    data = []
    for i, hit in enumerate(results[0]):
        data.append([hit.pk_id, hit.score, hit.video_time, hit.video_context])

    # 打印表格
    print(tabulate(data, headers=headers))

    # 释放加载的集合，减少内存消耗
    # conn.release()

    # 断开与 Milvus 服务器的连接
    # conn.disconnect()

    return results[0]


if __name__ == '__main__':
    # 文本数据
    text_data = """
    00:00:01~00:00:07
    欢迎来到Onboard,真实的一线经验,走新的投资思考。我是Monica。


    00:00:07~00:00:11
    我是高宁。我们一起聊聊软件如何改变世界。


    00:00:15~00:00:17
    大家好,欢迎来到Onboard,我是Monica。


    00:00:17~00:00:28
    自从OpenAI发布的ChatGBT掀起了席卷世界的AI热潮,不到三个月就积累了超过一亿的越活用户,超过1300万的日活用户。


    00:00:28~00:00:34
    真的是展现了AI让人惊叹的能力,也让很多人直呼这就是下一个互联网的未来。


    00:00:34~00:00:41
    有不少观众都说希望我们再做一期AI的讨论,于是这次硬核讨论就来了。


    00:00:41~00:00:49
    这次我们请来了Google Brain的研究员雪芝,她是Google大语言模型PALM Pathways Language Model的作者之一。


    00:00:49~00:00:53
    要知道,这个模型的参数量是GPT-3的三倍还多。


    00:00:53~00:01:01
    另外还有两位AI产品大牛,一位来自著名的Stable Diffusion背后的商业公司Surbility AI,


    00:01:01~00:01:08
    另一位来自某硅谷科技大厂,也曾在吴恩达教授的Landing AI中担任产品负责人。


    00:01:08~00:01:15
    此外,Monica还邀请到一位一直关注AI的投资人朋友Bill当作我的特邀共同主持嘉宾。


    00:01:15~00:01:22
    我们主要讨论几个话题,一方面从研究的视角,最前沿的研究者在关注什么?


    00:01:22~00:01:26
    现在的技术的天花板和未来大的变量可能会在哪里?


    00:01:26~00:01:30
    从产品和商业的角度,什么是一个好的AI产品?


    00:01:30~00:01:33
    整个生态可能随着技术有怎样的演变?


    00:01:33~00:01:38
    更重要的,我们又能从上一波AI的创业热潮中学到什么?


    00:01:38~00:01:44
    最后,Monica和Bill还会从投资人的视角做一个回顾、总结和畅想。


    00:01:44~00:01:51
    这里还有一个小的update,在本集发布的时候,Google也对爆发式增长的Chat GPT做出了回应。


    00:01:51~00:01:56
    正在测试一个基于Lambda模型的聊天机器人ApprenticeBot。


    00:01:56~00:02:00
    证实发布后会有怎样的惊喜,我们都拭目以待。


    00:02:00~00:02:04
    AI无疑是未来几年最令人兴奋的变量之一。


    00:02:04~00:02:10
    Monica也希望未来能邀请到更多一线从业者从不同角度讨论这个话题。


    00:02:10~00:02:15
    不论是想要做创业、研究、产品还是投资的同学,


    00:02:15~00:02:27
    希望这些对话对于大家了解这些技术演进、商业的可能,甚至未来对于我们每个人、每个社会意味着什么都能引发一些思考,提供一些启发。


    00:02:27~00:02:34
    这次的讨论有些技术硬核,需要各位对生成式AI大模型都有一些基础了解。


    00:02:34~00:02:41
    讨论中涉及到的论文和重要概念也会总结在本集的简介中,供大家复习参考。


    00:02:41~00:02:47
    几位嘉宾在北美工作生活多年,夹杂英文在所难免,也请大家体谅了。


    00:02:47~00:03:12
    欢迎来到未来,大家enjoy!
    """
    # 创建数据表与向量数据
    create_milvus(text_data)

    # 根据视频中的旁白查询视频时间
    embeddings = model.encode("Monica还邀请到一位一直关注AI的投资人朋友Bill")
    # print(embeddings.tolist()) # 查看所有向量
    search_milvus(embeddings)
